#include <ros/ros.h>
#include <std_msgs/UInt8.h>
#include <sensor_msgs/BatteryState.h>
#include <geometry_msgs/Twist.h>
#include <geometry_msgs/TwistStamped.h>
#include <nav_msgs/Odometry.h>
#include <nav_msgs/Path.h>
#include <diagnostic_msgs/DiagnosticArray.h>
#include <diagnostic_msgs/DiagnosticStatus.h>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2_ros/transform_broadcaster.h>
#include <geometry_msgs/TransformStamped.h>
#include <cmath>
#include <memory>
#include <string>

#include "ugv_sdk/scout/scout_base.hpp"

using namespace westonrobot;

class ScoutMiniRosNode {
public:
  explicit ScoutMiniRosNode(ros::NodeHandle& nh, ros::NodeHandle& pnh)
      : nh_(nh), pnh_(pnh) {
    pnh_.param<std::string>("can_device", can_device_, std::string("can0"));
    pnh_.param<std::string>("odom_frame", odom_frame_, std::string("odom"));
    pnh_.param<std::string>("base_frame", base_frame_, std::string("base_link"));
    pnh_.param<double>("wheelbase", wheelbase_, 0.452);  // from manual
    pnh_.param<double>("track", track_, 0.450);
    pnh_.param<double>("max_linear", max_linear_, 1.5);
    pnh_.param<double>("max_angular", max_angular_, 1.5);

    battery_pub_ = nh_.advertise<sensor_msgs::BatteryState>("scout/battery", 10);
    motor_pub_ = nh_.advertise<diagnostic_msgs::DiagnosticArray>("scout/motor_states", 10);
    twist_pub_ = nh_.advertise<geometry_msgs::TwistStamped>("scout/motion", 20);
    odom_pub_ = nh_.advertise<nav_msgs::Odometry>("odom", 20);
    path_pub_ = nh_.advertise<nav_msgs::Path>("scout/path", 5);
    light_cmd_sub_ = nh_.subscribe("scout/light_cmd", 10, &ScoutMiniRosNode::LightCmdCb, this);
    cmd_vel_sub_ = nh_.subscribe("cmd_vel", 20, &ScoutMiniRosNode::CmdVelCb, this);

    ConnectSdk();
  }

  void SpinOnce() {
    if (!scout_) return;

    robot_state_ = scout_->GetScoutState();

    PublishBattery(robot_state_);
    PublishMotor(robot_state_);
    PublishMotion(robot_state_);
    IntegrateAndPublishOdom(robot_state_);
  }

private:
  void ConnectSdk() {
    // For this ugv_sdk layout, use ScoutBase API, assuming mini platform
    scout_.reset(new ScoutBase(true));
    scout_->Connect(can_device_);
    scout_->EnableCommandedMode();
  }

  void PublishBattery(const ScoutState& st) {
    sensor_msgs::BatteryState msg;
    msg.header.stamp = ros::Time::now();
    msg.voltage = st.battery_voltage;
    msg.percentage = NAN; // percentage not available in this SDK layout
    msg.power_supply_status = sensor_msgs::BatteryState::POWER_SUPPLY_STATUS_DISCHARGING;
    battery_pub_.publish(msg);
  }

  void PublishMotor(const ScoutState& st) {
    diagnostic_msgs::DiagnosticArray arr;
    arr.header.stamp = ros::Time::now();
    for (size_t i = 0; i < ScoutState::motor_num; ++i) {
      const auto& a = st.actuator_states[i];
      diagnostic_msgs::DiagnosticStatus d;
      d.level = diagnostic_msgs::DiagnosticStatus::OK;
      d.name = std::string("motor_") + std::to_string(i);
      d.message = "ok";
      d.values.reserve(4);
      d.values.emplace_back(); d.values.back().key = "rpm"; d.values.back().value = std::to_string(a.motor_rpm);
      d.values.emplace_back(); d.values.back().key = "current"; d.values.back().value = std::to_string(a.motor_current);
      d.values.emplace_back(); d.values.back().key = "temperature"; d.values.back().value = std::to_string(a.motor_temperature);
      d.values.emplace_back(); d.values.back().key = "driver_state"; d.values.back().value = std::to_string(a.driver_state);
      arr.status.push_back(d);
    }
    motor_pub_.publish(arr);
  }

  void PublishMotion(const ScoutState& st) {
    geometry_msgs::TwistStamped tw;
    tw.header.stamp = ros::Time::now();
    tw.twist.linear.x = st.linear_velocity;
    tw.twist.angular.z = st.angular_velocity;
    twist_pub_.publish(tw);
    last_vx_ = tw.twist.linear.x;
    last_wz_ = tw.twist.angular.z;
  }

  void IntegrateAndPublishOdom(const ScoutState&) {
    ros::Time now = ros::Time::now();
    double dt = (now - last_time_).toSec();
    if (last_time_.toSec() == 0.0) {
      last_time_ = now;
      return;
    }
    // simple unicycle integration
    double delta_yaw = last_wz_ * dt;
    yaw_ += delta_yaw;
    double dx = last_vx_ * std::cos(yaw_) * dt;
    double dy = last_vx_ * std::sin(yaw_) * dt;
    x_ += dx;
    y_ += dy;

    geometry_msgs::TransformStamped tf;
    tf.header.stamp = now;
    tf.header.frame_id = odom_frame_;
    tf.child_frame_id = base_frame_;
    tf.transform.translation.x = x_;
    tf.transform.translation.y = y_;
    tf.transform.translation.z = 0.0;
    tf2::Quaternion q;
    q.setRPY(0, 0, yaw_);
    tf.transform.rotation.x = q.x();
    tf.transform.rotation.y = q.y();
    tf.transform.rotation.z = q.z();
    tf.transform.rotation.w = q.w();
    tf_broadcaster_.sendTransform(tf);

    nav_msgs::Odometry odom;
    odom.header.stamp = now;
    odom.header.frame_id = odom_frame_;
    odom.child_frame_id = base_frame_;
    odom.pose.pose.position.x = x_;
    odom.pose.pose.position.y = y_;
    odom.pose.pose.position.z = 0.0;
    odom.pose.pose.orientation.x = q.x();
    odom.pose.pose.orientation.y = q.y();
    odom.pose.pose.orientation.z = q.z();
    odom.pose.pose.orientation.w = q.w();
    odom.twist.twist.linear.x = last_vx_;
    odom.twist.twist.angular.z = last_wz_;
    odom_pub_.publish(odom);

    // publish smooth path without arrows
    geometry_msgs::PoseStamped ps;
    ps.header.stamp = now;
    ps.header.frame_id = odom_frame_;
    ps.pose.position.x = x_;
    ps.pose.position.y = y_;
    ps.pose.position.z = 0.0;
    ps.pose.orientation = odom.pose.pose.orientation;
    path_msg_.header.stamp = now;
    path_msg_.header.frame_id = odom_frame_;
    path_msg_.poses.push_back(ps);
    if (path_msg_.poses.size() > 1000) {
      path_msg_.poses.erase(path_msg_.poses.begin(), path_msg_.poses.begin() + 200);
    }
    path_pub_.publish(path_msg_);

    last_time_ = now;
  }

  void CmdVelCb(const geometry_msgs::Twist::ConstPtr& msg) {
    if (!scout_) return;
    double v = std::max(std::min((double)msg->linear.x, max_linear_), -max_linear_);
    double w = std::max(std::min((double)msg->angular.z, max_angular_), -max_angular_);
    scout_->SetMotionCommand(v, w);
  }

  void LightCmdCb(const std_msgs::UInt8::ConstPtr& msg) {
    if (!scout_) {
      ROS_WARN("scout_mini_node: Light command received, but SDK not connected.");
      return;
    }

    ROS_INFO("scout_mini_node: Received light command: %d", msg->data);

    ScoutLightCmd cmd;
    bool disable_light = false;

    switch (msg->data) {
      case 0: // off
        ROS_INFO("  -> Action: Turn OFF lights.");
        disable_light = true;
        break;
      case 1: // constant on
        ROS_INFO("  -> Action: Set lights to CONSTANT ON.");
        cmd = ScoutLightCmd(ScoutLightCmd::LightMode::CONST_ON, 0,
                            ScoutLightCmd::LightMode::CONST_ON, 0);
        break;
      case 2: // breath
        ROS_INFO("  -> Action: Set lights to BREATH mode.");
        cmd = ScoutLightCmd(ScoutLightCmd::LightMode::BREATH, 0,
                            ScoutLightCmd::LightMode::BREATH, 0);
        break;
      case 3: // custom blink-ish
        ROS_INFO("  -> Action: Set lights to CUSTOM mode.");
        cmd = ScoutLightCmd(ScoutLightCmd::LightMode::CUSTOM, 10,
                            ScoutLightCmd::LightMode::CUSTOM, 10);
        break;
      default:
        ROS_WARN("  -> Action: Unknown command (%d), turning lights OFF.", msg->data);
        disable_light = true;
        break;
    }

    if (disable_light) {
        scout_->DisableLightCmdControl();
    } else {
        scout_->SetLightCommand(cmd);
    }
}

private:
  ros::NodeHandle nh_;
  ros::NodeHandle pnh_;

  std::string can_device_;
  std::string odom_frame_;
  std::string base_frame_;
  double wheelbase_{0.452};
  double track_{0.450};
  double max_linear_{1.5};
  double max_angular_{1.5};

  ros::Publisher battery_pub_;
  ros::Publisher motor_pub_;
  ros::Publisher twist_pub_;
  ros::Publisher odom_pub_;
  ros::Publisher path_pub_;
  ros::Subscriber cmd_vel_sub_;
  ros::Subscriber light_cmd_sub_;

  std::unique_ptr<ScoutBase> scout_;
  ScoutState robot_state_;

  tf2_ros::TransformBroadcaster tf_broadcaster_;
  ros::Time last_time_;
  double x_{0.0}, y_{0.0}, yaw_{0.0};
  double last_vx_{0.0}, last_wz_{0.0};
  nav_msgs::Path path_msg_;
};

int main(int argc, char** argv) {
  ros::init(argc, argv, "scout_mini_node");
  ros::NodeHandle nh;
  ros::NodeHandle pnh("~");
  ScoutMiniRosNode node(nh, pnh);
  ros::Rate rate(50);
  while (ros::ok()) {
    node.SpinOnce();
    ros::spinOnce();
    rate.sleep();
  }
  return 0;
}



